import paho.mqtt.client as mqtt
import os
import logging
import time
import json
import subprocess
import requests
import random
import psutil
import threading
from publib import setup_logger, get_logger

cur_file = os.path.basename(__file__).split(".")[0]
logger = get_logger(__name__)

config = {}
if (os.path.exists("config.json")):
    with open("config.json", "r") as f:
        config = json.load(f)

device_name = config.get("device_name", "public")
mqtt_client = None
device = f"{device_name}_{cur_file}"
listen_devices = [device]
topic_type = ["up_reply", "down"]
sub_topic = [f"/xx/{dev}/{type}" for dev in listen_devices for type in topic_type]
msg_id = random.randint(1, 100000)
run_table = {}

def basic_cmd(msg):
    id = msg["msg_id"]
    cmd = msg["cmd"]
    pub_response({"cmd": cmd, "msg_id": id})

def exec_file(msg):
    id = msg["msg_id"]
    cmd = msg["cmd"]
    pub_response({"cmd": cmd, "msg_id": id})

    param = msg.get("param", {})
    path = param.get("path", "")
    name = path.split('/')[-1]
    content = param.get("content", None)
    url = param.get("url", None)
    if (url):
        try:
            content = requests.get(url, timeout=20).text
        except:
            content = None
    if (content):
        with open(path, "w", encoding='utf-8') as f:
            f.write(content)
    
    local_run(name, param.get("param", []))

def kill(msg):
    id = msg["msg_id"]
    cmd = msg["cmd"]
    pub_response({"cmd": cmd, "msg_id": id})

    name = msg.get("param", {}).get("name", "")
    local_kill(name)


cmd_table = {
    "heartbeats": basic_cmd,
    "online": basic_cmd,
    "exec_file": exec_file,
    "kill": kill,
}

def dl_file(url, path):
    logger.info(f"download {url} -> {path}")
    try:
        r = requests.get(url, timeout=20)
        with open(path, 'wb') as f:
            f.write(r.content)
    except:
        logger.info(f"download {url} failed")

def local_run(name, param):
    global run_table

    # kill要运行的进程
    local_kill(name)

    # 运行进程
    try:
        env = os.environ.copy()
        env['PYTHONIOENCODING'] = 'utf-8'
        with open('logs/stdout.log', 'w', encoding='utf-8') as stdout_file:
            process = subprocess.Popen(param, stdout=stdout_file, stderr=stdout_file, text=True, encoding="utf-8", env=env)
            logger.info(f"run {name}, pid: {process.pid}")
            run_table.update({name: {"pid": process.pid, "handle": process}})
    except:
        logger.info(f"run {name} failed")

def local_kill(name):
    def _kill_processes(pid=None):
        # kill当前进程及子进程
        if pid is not None:
            try:
                parent = psutil.Process(pid)
            except psutil.NoSuchProcess:
                logger.info(f"kill failed, {pid} not exists")
                return False
            try:
                children = parent.children(recursive=True)
                for proc in children + [parent]:
                    proc.terminate()
                gone, still_alive = psutil.wait_procs(children + [parent], timeout=5)
                for p in still_alive:
                    p.kill()
                print(f"kill {pid} success")
                return True
            except Exception as e:
                print(f"kill {pid} failed: {e}")
                return False
    
    global run_table
    # kill要运行的进程
    if (name in run_table):
        logger.info(f"kill {name}")
        _kill_processes(run_table[name]["handle"].pid)
        run_table.pop(name)

def _pub_msg(topic, msg):
    if (mqtt_client and mqtt_client.is_connected() == False):
        logger.info("pub err, mqtt not connected")
        return
    if (isinstance(msg, dict)):
        global msg_id
        msg_id += 1
        if ("msg_id" not in msg):
            msg.update({"msg_id": msg_id})
        msg = json.dumps(msg)
    mqtt_client.publish(topic, msg, qos=2)
    if ("heartbeats" not in msg):
        logger.info(f"pub {topic} {msg}")

def pub_cmd(msg):
    _pub_msg(f"/xx/{device}/up", msg)

def pub_response(msg):
    _pub_msg(f"/xx/{device}/down_reply", msg)

def cb_recv_msg(topic, msg):
    type = topic.split('/')[-1]
    if (type == "down"):
        cmd = msg["cmd"]
        if (cmd in cmd_table):
            cmd_table[cmd](msg)

def mqtt_start():
    def on_connect(client, userdata, flags, rc):
        logger.info(f"Connected with result code: {rc}")
        if (rc != 0):
            return
        
        for topic in sub_topic:
            mqtt_client.subscribe(topic, qos=2)
        
        msg = {"cmd": "online"}
        pub_cmd(msg)

    def on_disconnect(client, userdata, rc):
        logger.info(f"Disconnected with result code: {rc}")

    def on_message(client, userdata, msg):
        logger.info(f"recv {msg.topic} {msg.payload.decode()}")
        try:
            cb_recv_msg(msg.topic, json.loads(msg.payload))
        except:
            logger.info(f"process msg error")

    # 连接mqtt服务器
    global mqtt_client
    mqtt_addr = config.get("mqtt_addr", "")
    if (mqtt_addr == ""):
        logger.error(f"mqtt_addr is empty")
        return
    mqtt_host = mqtt_addr.split(":")[0]
    mqtt_port = int(mqtt_addr.split(":")[1])
    mqtt_client = mqtt.Client(client_id = device)
    mqtt_client.reconnect_delay_set(min_delay=5, max_delay=600)
    mqtt_client.on_connect = on_connect
    mqtt_client.on_message = on_message
    mqtt_client.on_disconnect = on_disconnect
    mqtt_client.connect(mqtt_host, mqtt_port, 60) # 600为keepalive的时间间隔
    mqtt_client.loop_start()
    mqtt_client.enable_logger(logging)
    logger.info(f"mqtt task start")

def main_thread():
    while True:
        time.sleep(60 * 30)
        msg = {"cmd": "heartbeats"}
        try:
            pub_cmd(msg)
        except:
            logger.error(f"pub heartbeats failed")

def main():
    setup_logger(f"logs/{cur_file}.log")
    if (cur_file == "main"):
        with open(f"{cur_file}.pid", "w") as f:
            f.write(f"{os.getpid()}")
        local_run("mqtt.py", ["python", "mqtt.py"])
    else:
        for p in config.get("run", []):
            local_run(p, config["run"][p])
    mqtt_start()

    # 创建进程
    main_thread_handle = threading.Thread(target=main_thread, daemon=True)
    main_thread_handle.start()

    try:
        while (1):
            time.sleep(3600)
    except KeyboardInterrupt:
        logger.info("KeyboardInterrupt, exit")


if __name__ == "__main__":
    main()
